#!/usr/bin/env python3
"""
MPC (QP) 版 + 简单避障实验（绕箱体上方）
-------------------------------------------------
亮点：
1) 用 cvxpy 解真正的二次规划 (QP)，把输入/状态约束直接写进优化问题；
2) 仍使用线性双积分模型 (pos-vel-acc)，但接口清晰，便于替换更复杂模型；
3) 引入“上方绕行”的线性半空间约束 y >= box.y_max + margin（当仍位于障碍前方时才施加）；
4) 结合速度受限的参考路径（waypoint）形成可执行走廊；
5) 保留实时可视化：3D 轨迹、MPC 预测（含速度箭头）、控制序列时间窗、u_raw(=优化解)；

依赖：
  pip install cvxpy osqp numpy pandas matplotlib

注意：
 - 为保证优化始终是凸的/快速的，避障采用“线性半空间约束”的工程化路线：
   在尚未越过障碍（x <= box.x_max+margin）时强制 y 高于障碍上沿 (y >= y_gate)。
   这是典型的走廊化处理，能稳定展示 MPC 在“有约束”的优势。
 - 如需更复杂的避障（任意绕行/多障碍/非凸避障），建议使用混合整数或顺序凸优化（SCP）。
"""
import os
import time as pytime
from math import ceil
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import cvxpy as cp

# =====================
# 参数区域（可按需调整）
# =====================
DT = 0.1
T_TOTAL = 60.0
STEPS = int(ceil(T_TOTAL / DT))

# 代价权重
Q_pos = 10.0
Q_vel = 2.0
R_acc = 0.1

# 终端权重（可略大，利于收敛）
Qf_pos = 25.0
Qf_vel = 5.0

# 地面/飞行空间约束（可选）
POS_MIN = np.array([-10.0, -30.0,  0.0])   # 例如 z>=0
POS_MAX = np.array([120.0,  30.0, 50.0])
VEL_MAX = 12.0

# 执行器约束（加速度）
U_MIN = np.array([-5.0, -5.0, -5.0])
U_MAX = np.array([ 5.0,  5.0,  5.0])

# 目标与初始
TARGET_POS = np.array([80.0, -20.0, 30.0])
X0 = np.zeros(6); X0[2] = 0.0

# 预测时域
HORIZON = 30
PRED_SHOW = 12

# 到达/悬停判据
POS_TOL = 0.5
VEL_TOL = 0.5
HOVER_TIME = 6.0

# Waypoint 速度限制（用于参考生成）
WAYPOINT_MAX_SPEED = 6.0

# 障碍物（轴对齐立方体）
BOX_MIN = np.array([35.0, -8.0,  0.0])
BOX_MAX = np.array([55.0,  8.0, 22.0])
SAFETY_MARGIN = 2.0
Y_GATE = BOX_MAX[1] + SAFETY_MARGIN   # 走廊：从上方绕行 -> y 必须大于此阈值

# 可视化开关
SHOW_INTERACTIVE = True
SAVE_FIG = True

# ===============
# 动力学（双积分）
# ===============
I3 = np.eye(3)
A = np.block([[I3, DT * I3],
              [np.zeros((3,3)), I3]])
B = np.block([[0.5 * DT**2 * I3],
              [DT * I3]])

Q = np.diag([Q_pos, Q_pos, Q_pos, Q_vel, Q_vel, Q_vel])
Qf = np.diag([Qf_pos, Qf_pos, Qf_pos, Qf_vel, Qf_vel, Qf_vel])
R = R_acc * np.eye(3)

# =====================
# 参考路径（速度受限）
# =====================

def generate_speed_limited_ref(x_start, target_pos, dt, max_speed, steps):
    start_pos = x_start[0:3].copy()
    vec = target_pos - start_pos
    dist = np.linalg.norm(vec)
    if dist < 1e-6:
        return np.tile(target_pos, (steps,1))
    dir_unit = vec / dist
    max_step = max_speed * dt
    n_steps_needed = int(np.ceil((dist - 0.5) / max_step)) if dist > 0.5 else 1
    n_steps_needed = max(1, n_steps_needed)
    pos_seq = [start_pos + dir_unit * min(max_step * i, dist) for i in range(n_steps_needed+1)]
    if len(pos_seq) < steps:
        pos_seq.extend([target_pos] * (steps - len(pos_seq)))
    return np.array(pos_seq[:steps])

ref_pos_seq = generate_speed_limited_ref(X0, TARGET_POS, DT, WAYPOINT_MAX_SPEED, STEPS)

def ref_state_at(step_idx):
    p = ref_pos_seq[min(step_idx, ref_pos_seq.shape[0]-1)]
    xr = np.zeros(6); xr[0:3] = p; xr[3:6] = 0.0
    return xr

# =====================
# QP MPC 求解器（cvxpy + OSQP）
# =====================

def solve_mpc_qp(x_curr, x_ref, horizon,
                 A, B, Q, R, Qf,
                 u_min, u_max,
                 pos_min=None, pos_max=None,
                 vel_max=None,
                 enforce_above=None):
    """
    enforce_above: 若给出一个标量 y_gate（float），则添加约束 Y_k >= y_gate （k=0..N）
                   用于实现“上方绕行”的线性安全走廊。
    """
    nx = A.shape[0]; nu = B.shape[1]
    X = cp.Variable((nx, horizon+1))
    U = cp.Variable((nu, horizon))

    cost = 0
    cons = []

    # 初始状态
    cons.append(X[:,0] == x_curr)

    # 动力学 & 约束 & 代价
    for k in range(horizon):
        # 动力学
        cons.append(X[:,k+1] == A @ X[:,k] + B @ U[:,k])
        # 输入约束
        cons.append(U[:,k] <= u_max)
        cons.append(U[:,k] >= u_min)
        # 状态约束（可选）
        if pos_min is not None:
            cons.append(X[0:3,k] >= pos_min)
        if pos_max is not None:
            cons.append(X[0:3,k] <= pos_max)
        if vel_max is not None:
            cons.append(cp.norm_inf(X[3:6,k]) <= vel_max)
        if enforce_above is not None:
            cons.append(X[1,k] >= float(enforce_above))  # y >= gate
        # stage cost
        x_err = X[:,k] - x_ref
        cost += cp.quad_form(x_err, Q) + cp.quad_form(U[:,k], R)

    # 末端约束
    if pos_min is not None:
        cons.append(X[0:3,horizon] >= pos_min)
    if pos_max is not None:
        cons.append(X[0:3,horizon] <= pos_max)
    if vel_max is not None:
        cons.append(cp.norm_inf(X[3:6,horizon]) <= vel_max)
    if enforce_above is not None:
        cons.append(X[1,horizon] >= float(enforce_above))

    # 末端代价
    x_errN = X[:,horizon] - x_ref
    cost += cp.quad_form(x_errN, Qf)

    prob = cp.Problem(cp.Minimize(cost), cons)
    prob.solve(solver=cp.OSQP, warm_start=True, verbose=False)

    if prob.status not in [cp.OPTIMAL, cp.OPTIMAL_INACCURATE]:
        raise RuntimeError(f"MPC QP infeasible or not optimal: {prob.status}")

    u0 = U[:,0].value
    X_pred = X.value
    U_pred = U.value
    return u0, X_pred, U_pred

# =============
# 悬停控制 PD
# =============

def pd_hover_control(pos, vel, target_pos, kp=0.6, kd=1.0):
    pos_err = pos - target_pos
    return -kp * pos_err - kd * vel

# ============
# 可视化工具
# ============

def setup_plot():
    fig = plt.figure(figsize=(14,8))
    gs = fig.add_gridspec(2, 3, width_ratios=[2.8,1,1])
    ax3d = fig.add_subplot(gs[:,0], projection='3d')
    ax_pos = fig.add_subplot(gs[0,1])
    ax_ctrl = fig.add_subplot(gs[1,1])
    ax_u = fig.add_subplot(gs[0,2])
    ax_info = fig.add_subplot(gs[1,2]); ax_info.axis('off')

    traj_line, = ax3d.plot([], [], [], lw=2, label='trajectory')
    start_scatter = ax3d.scatter([X0[0]],[X0[1]],[X0[2]], c='b', s=80, marker='o', label='start')
    target_scatter = ax3d.scatter([TARGET_POS[0]],[TARGET_POS[1]],[TARGET_POS[2]], c='r', s=120, marker='x', label='target')
    box_draw(ax3d, BOX_MIN, BOX_MAX, color='tab:gray', alpha=0.2)

    pred_line, = ax3d.plot([], [], [], ls='--', lw=1.3, color='orange', label='pred_traj')
    pred_quiver = None
    final_scatter = ax3d.scatter([], [], [], c='g', s=150, marker='*', label='final')

    ax3d.set_xlabel('x (m)'); ax3d.set_ylabel('y (m)'); ax3d.set_zlabel('z (m)')
    ax3d.set_title('3D Trajectory (QP-MPC with obstacle corridor)')
    ax3d.legend(loc='upper left')

    pos_lines = [ax_pos.plot([], [], label='x')[0],
                 ax_pos.plot([], [], label='y')[0],
                 ax_pos.plot([], [], label='z')[0]]
    ax_pos.set_xlabel('time (s)'); ax_pos.set_ylabel('position (m)'); ax_pos.set_title('Position vs time'); ax_pos.legend(); ax_pos.grid(True)

    u_lines = [ax_ctrl.plot([], [], label='ax')[0],
               ax_ctrl.plot([], [], label='ay')[0],
               ax_ctrl.plot([], [], label='az')[0]]
    ax_ctrl.set_xlabel('time (s)'); ax_ctrl.set_ylabel('acc (m/s^2)'); ax_ctrl.set_title('Applied control u'); ax_ctrl.legend(); ax_ctrl.grid(True)

    uraw_lines = [ax_u.plot([], [], label='u_raw_x')[0],
                  ax_u.plot([], [], label='u_raw_y')[0],
                  ax_u.plot([], [], label='u_raw_z')[0]]
    ax_u.set_title('u_raw (QP solution)'); ax_u.set_xlabel('time (s)'); ax_u.set_ylabel('acc (m/s^2)'); ax_u.legend(); ax_u.grid(True)

    info_text = ax_info.text(0.01, 0.98, "", va='top', ha='left', fontsize=9, family='monospace')

    return fig, ax3d, ax_pos, ax_ctrl, ax_u, info_text, traj_line, pred_line, final_scatter, pos_lines, u_lines, uraw_lines


def box_draw(ax, bmin, bmax, color='k', alpha=0.15):
    # 画轴对齐盒子（线框 + 半透明 patch）
    xs = [bmin[0], bmax[0]]; ys = [bmin[1], bmax[1]]; zs = [bmin[2], bmax[2]]
    for x in xs:
        for y in ys:
            ax.plot([x,x], [y,y], [zs[0], zs[1]], color=color, alpha=0.6)
    for x in xs:
        for z in zs:
            ax.plot([x,x], [ys[0], ys[1]], [z,z], color=color, alpha=0.6)
    for y in ys:
        for z in zs:
            ax.plot([xs[0], xs[1]], [y,y], [z,z], color=color, alpha=0.6)
    # 顶面半透明（简单展示）
    Xp = [xs[0], xs[1], xs[1], xs[0], xs[0]]
    Yp = [ys[0], ys[0], ys[1], ys[1], ys[0]]
    Zp = [zs[1]]*5
    ax.plot(Xp, Yp, Zp, color=color, alpha=alpha)


def update_plots(fig, ax3d, ax_pos, ax_ctrl, ax_u,
                 info_text,
                 traj_line, pred_line, final_scatter,
                 pos_lines, u_lines, uraw_lines,
                 t_arr, X_arr, U_arr, Uraw_arr,
                 X_pred=None, U_pred=None):
    traj_line.set_data(X_arr[:,0], X_arr[:,1]); traj_line.set_3d_properties(X_arr[:,2])

    if X_pred is not None and X_pred.shape[1] > 1:
        pred_line.set_data(X_pred[0,:], X_pred[1,:]); pred_line.set_3d_properties(X_pred[2,:])
        # 速度箭头
        vel = X_pred[3:6,:]
        scale = 0.2
        Xp, Yp, Zp = X_pred[0,:], X_pred[1,:], X_pred[2,:]
        Up, Vp, Wp = vel[0,:]*scale, vel[1,:]*scale, vel[2,:]*scale
        try:
            ax3d.quiver(Xp, Yp, Zp, Up, Vp, Wp, length=1.0, normalize=False, color='orange', linewidth=1)
        except Exception:
            for i in range(Xp.shape[0]):
                ax3d.plot([Xp[i], Xp[i]+Up[i]], [Yp[i], Yp[i]+Vp[i]], [Zp[i], Zp[i]+Wp[i]], color='orange', alpha=0.8)
    else:
        pred_line.set_data([], []); pred_line.set_3d_properties([])

    # 时域轴范围
    if t_arr.size == 1:
        t0 = float(t_arr[0])
        for ax in (ax_pos, ax_ctrl, ax_u):
            ax.set_xlim(t0-0.1, t0+0.1)
    else:
        for ax in (ax_pos, ax_ctrl, ax_u):
            ax.set_xlim(t_arr[0], t_arr[-1])

    # pos / control
    for i, ln in enumerate(pos_lines):
        ln.set_data(t_arr, X_arr[:, i])
    all_pos = X_arr[:,0:3].flatten(); ax_pos.set_ylim(np.min(all_pos)-5, np.max(all_pos)+5)

    for i, ln in enumerate(u_lines):
        ln.set_data(t_arr, U_arr[:, i])
    if U_arr.size>0:
        all_u = U_arr.flatten(); ax_ctrl.set_ylim(np.min(all_u)-1.0, np.max(all_u)+1.0)

    for i, ln in enumerate(uraw_lines):
        ln.set_data(t_arr, Uraw_arr[:, i])
    if Uraw_arr.size>0:
        all_ur = Uraw_arr.flatten(); ax_u.set_ylim(np.min(all_ur)-1.0, np.max(all_ur)+1.0)

    fig.canvas.draw(); fig.canvas.flush_events()

# ============
# 主循环
# ============

def main():
    out_dir = os.path.dirname(os.path.abspath(__file__))
    log_path = os.path.join(out_dir, 'mpc_qp_obstacle_log.csv')
    fig_path = os.path.join(out_dir, 'mpc_qp_obstacle.png')

    if SHOW_INTERACTIVE:
        plt.ion()
        (fig, ax3d, ax_pos, ax_ctrl, ax_u,
         info_text,
         traj_line, pred_line, final_scatter,
         pos_lines, u_lines, uraw_lines) = setup_plot()
    else:
        plt.ioff(); fig=None

    times, X_hist, U_hist, Uraw_hist = [], [], [], []

    x = X0.copy()
    hover_mode = False; hover_start = None
    reached = False

    sim_start = pytime.time()

    for step in range(STEPS):
        t = step * DT
        xr = ref_state_at(step)

        # 是否启用上方绕行约束：尚未越过障碍的 x 范围就启用
        enforce_gate = None
        if x[0] <= BOX_MAX[0] + SAFETY_MARGIN:  # 仍在障碍前方
            enforce_gate = Y_GATE

        # MPC 求解
        if not hover_mode:
            try:
                u_opt, Xp, Up = solve_mpc_qp(
                    x_curr=x, x_ref=xr, horizon=HORIZON,
                    A=A, B=B, Q=Q, R=R, Qf=Qf,
                    u_min=U_MIN, u_max=U_MAX,
                    pos_min=POS_MIN, pos_max=POS_MAX,
                    vel_max=VEL_MAX,
                    enforce_above=enforce_gate,
                )
            except RuntimeError as e:
                print("QP infeasible, fallback to PD hover:", e)
                u_opt = pd_hover_control(x[0:3], x[3:6], TARGET_POS)
                Xp, Up = None, None
        else:
            u_opt = pd_hover_control(x[0:3], x[3:6], TARGET_POS)
            Xp, Up = None, None

        # 应用输入，并记录（这里 u_raw = u_opt）
        u = np.clip(u_opt, U_MIN, U_MAX)
        w = (np.random.randn(3)*0.0)  # 可加入微扰
        x = A.dot(x) + B.dot(u) + B.dot(w)

        times.append(t)
        X_hist.append(x.copy())
        U_hist.append(u.copy())
        Uraw_hist.append(u_opt.copy())

        # 到达 + 悬停
        pos = x[0:3]; vel = x[3:6]
        err_to_target = np.linalg.norm(pos - TARGET_POS)
        if (err_to_target < POS_TOL) and (np.linalg.norm(vel) < VEL_TOL) and (not hover_mode):
            hover_mode = True; hover_start = t
            print(f"*** ARRIVED at FINAL TARGET -> HOVER {HOVER_TIME}s (t={t:.2f}) ***")

        if hover_mode and (t - hover_start >= HOVER_TIME):
            print("*** HOVER complete -> ending sim ***"); reached = True; break

        # 可视化
        try:
            if SHOW_INTERACTIVE and fig is not None:
                t_arr = np.array(times)
                X_arr = np.array(X_hist)
                U_arr = np.array(U_hist)
                Uraw_arr = np.array(Uraw_hist)
                update_plots(fig, ax3d, ax_pos, ax_ctrl, ax_u,
                             info_text,
                             traj_line, pred_line, final_scatter,
                             pos_lines, u_lines, uraw_lines,
                             t_arr, X_arr, U_arr, Uraw_arr,
                             X_pred=Xp, U_pred=Up)
                plt.pause(0.001)
        except Exception as e:
            print("Plot update error:", e)
            break

    sim_end = pytime.time()
    print(f"Simulation ended. wall_time={sim_end-sim_start:.2f}s steps_simulated={len(times)} reached={reached}")

    # 保存日志与图
    Xnp = np.array(X_hist); Unp = np.array(U_hist); Urawnp = np.array(Uraw_hist)
    df = pd.DataFrame({
        't': np.array(times),
        'x': Xnp[:,0], 'y': Xnp[:,1], 'z': Xnp[:,2],
        'vx': Xnp[:,3], 'vy': Xnp[:,4], 'vz': Xnp[:,5],
        'ax': Unp[:,0], 'ay': Unp[:,1], 'az': Unp[:,2],
        'u_raw_x': Urawnp[:,0], 'u_raw_y': Urawnp[:,1], 'u_raw_z': Urawnp[:,2],
    })
    df.to_csv(log_path, index=False)
    print("Saved CSV log to:", log_path)

    try:
        if SAVE_FIG:
            if SHOW_INTERACTIVE and fig is not None:
                fig.savefig(fig_path, dpi=160, bbox_inches='tight')
            else:
                # 生成静态图
                fig_s = plt.figure(figsize=(14,8))
                gs2 = fig_s.add_gridspec(2,3, width_ratios=[2.8,1,1])
                a3 = fig_s.add_subplot(gs2[:,0], projection='3d')
                ap = fig_s.add_subplot(gs2[0,1]); ac = fig_s.add_subplot(gs2[1,1]); au = fig_s.add_subplot(gs2[0,2])
                a3.plot(Xnp[:,0], Xnp[:,1], Xnp[:,2], lw=2)
                a3.scatter([X0[0]],[X0[1]],[X0[2]], c='b', s=80, marker='o', label='start')
                a3.scatter([TARGET_POS[0]],[TARGET_POS[1]],[TARGET_POS[2]], c='r', s=120, marker='x', label='target')
                box_draw(a3, BOX_MIN, BOX_MAX, color='tab:gray', alpha=0.2)
                a3.set_title('3D Trajectory (final)'); a3.set_xlabel('x'); a3.set_ylabel('y'); a3.set_zlabel('z')
                ap.plot(df['t'], df['x'], label='x'); ap.plot(df['t'], df['y'], label='y'); ap.plot(df['t'], df['z'], label='z'); ap.legend(); ap.grid(True)
                ac.plot(df['t'], df['ax'], label='ax'); ac.plot(df['t'], df['ay'], label='ay'); ac.plot(df['t'], df['az'], label='az'); ac.legend(); ac.grid(True)
                au.plot(df['t'], df['u_raw_x'], label='u_raw_x'); au.plot(df['t'], df['u_raw_y'], label='u_raw_y'); au.plot(df['t'], df['u_raw_z'], label='u_raw_z'); au.legend(); au.grid(True)
                fig_s.tight_layout(); fig_s.savefig(fig_path, dpi=160, bbox_inches='tight')
            print("Saved figure to:", fig_path)
    except Exception as e:
        print("Failed to save figure:", e)

    if SHOW_INTERACTIVE:
        try:
            print("Interactive: close figure window to finish.")
            plt.ioff(); plt.show()
        except Exception:
            pass


if __name__ == "__main__":
    main()
